/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: GPU/NPU send/recv mr manager
 */
#include "typical_mr_manager.h"
#include "adapter_rts_common.h"
#include "adapter_hccp_common.h"
#include "adapter_rts.h"
#include "network_manager_pub.h"
#include "rdma_resource_manager.h"

namespace hccl {
constexpr uint32_t TypicalMrManager::DEFAULT_MR_KEY;

TypicalMrManager &TypicalMrManager::GetInstance()
{
    static TypicalMrManager typicalMrManager;
    return typicalMrManager;
}

TypicalMrManager::TypicalMrManager()
    : rdmaHandle_(nullptr)
{
}

TypicalMrManager::~TypicalMrManager()
{
    ReleaseMrResource();
}

HcclResult TypicalMrManager::RegisterMem(struct mr_info &mrInfo)
{
    HCCL_DEBUG("[TypicalMrManager][RegisterMem]MR register start, addr[%llu], MR size[%llu].",
        mrInfo.addr, mrInfo.size);
    
    if (!rdmaHandle_) {
        CHK_RET(RdmaResourceManager::GetInstance().GetRdmaHandle(rdmaHandle_));
    }
    CHK_PTR_NULL(rdmaHandle_);

    // Check whether the key is the default value or already registered in map
    std::unique_lock<std::mutex> lockMrMap(mrMapMutex_);
    if (mrInfo.lkey != DEFAULT_MR_KEY) {
        auto mrIter = regedMrMap_.find(mrInfo.lkey);
        if (mrIter != regedMrMap_.end()) {
            HCCL_WARNING("[TypicalMrManager][RegisterMem]MR key[%lu] already registered, " \
                "MR size[%llu], regedMrMap size[%u].",
                mrInfo.lkey, mrIter->second.first.size, regedMrMap_.size());
            return HCCL_E_PARA;
        }
        HCCL_ERROR("[TypicalMrManager][RegisterMem]invalid MR info for register, addr[%llu], key[%lu].",
            mrInfo.addr, mrInfo.lkey);
        return HCCL_E_PARA;
    }
    lockMrMap.unlock();

    MrHandle mrHandle = nullptr;
    // Register MR
    CHK_RET(hrtRaRegGlobalMr(rdmaHandle_, mrInfo, mrHandle));
    if (mrHandle == nullptr) {
        HCCL_WARNING("[TypicalMrManager][RegisterMem]MR register not success, addr[%llu], MR size[%llu].",
            mrInfo.addr, mrInfo.size);
        return HCCL_E_INTERNAL;
    }

    // Add the registered MR into MR map
    lockMrMap.lock();
    regedMrMap_[mrInfo.lkey].first = mrInfo;
    regedMrMap_[mrInfo.lkey].second = mrHandle;
    lockMrMap.unlock();

    HCCL_INFO("[TypicalMrManager][RegisterMem]MR register success, " \
        "MR key[%llu], addr[%llu], size[%llu], mrHandle[%p], regedMrMap size[%u].",
        mrInfo.lkey, mrInfo.addr, mrInfo.size, mrHandle, regedMrMap_.size());
    return HCCL_SUCCESS;
}

HcclResult TypicalMrManager::DeRegisterMem(struct mr_info &mrInfo)
{
    HCCL_DEBUG("[TypicalMrManager][DeRegisterMem]MR deregister start, addr[%llu], size[%llu], key[%lu].",
        mrInfo.addr, mrInfo.size, mrInfo.lkey);
    CHK_PTR_NULL(rdmaHandle_);
    // Check whether the mem key exists. If exists，remove from MR map; else return error.
    std::unique_lock<std::mutex> lockMrMap(mrMapMutex_);
    auto mrIter = regedMrMap_.find(mrInfo.lkey);
    if (mrIter == regedMrMap_.end()) {
        HCCL_ERROR("[TypicalMrManager][DeRegisterMem]no match MR info in MR map, " \
            "MR key[%llu], addr[%llu], size[%llu], regedMrMap size[%u].",
            mrInfo.lkey, mrInfo.addr, mrInfo.size, regedMrMap_.size());
        return HCCL_E_PARA;
    }
    lockMrMap.unlock();

    // Unregister MR
    MrHandle mrHandle = mrIter->second.second;
    HCCL_DEBUG("[TypicalMrManager][DeRegisterMem]rdma handle[%p], mr handle [%p].", rdmaHandle_, mrHandle);
    CHK_RET(hrtRaDeRegGlobalMr(rdmaHandle_, mrHandle));

    // Remove unregistered MR from MR map
    lockMrMap.lock();
    regedMrMap_.erase(mrIter);
    lockMrMap.unlock();
    mrInfo.lkey = DEFAULT_MR_KEY;

    HCCL_INFO("[TypicalMrManager][DeRegisterMem]MR unregister success, regedMrMap size[%u].",
        regedMrMap_.size());
    return HCCL_SUCCESS;
}

HcclResult TypicalMrManager::ReleaseMrResource()
{
    CHK_PTR_NULL(rdmaHandle_);
    std::unique_lock<std::mutex> lockMrMap(mrMapMutex_);
    if (!regedMrMap_.empty()) {
        for (auto &mrIter : regedMrMap_) {
            MrHandle mrHandle = mrIter.second.second;
            if (mrHandle != nullptr) {
                CHK_RET(hrtRaDeRegGlobalMr(rdmaHandle_, mrHandle));
            }
        }
        regedMrMap_.clear();
    }
    HCCL_INFO("[TypicalMrManager][ReleaseMrResource]release mr resources success.");
    return HCCL_SUCCESS;
}
}  // namespace hccl
